package cn.fengtang.chatglm.data.test;

import cn.fengtang.chatglm.model.ChatCompletionRequest;
import cn.fengtang.chatglm.model.ChatCompletionResponse;
import cn.fengtang.chatglm.model.Model;
import cn.fengtang.chatglm.model.Role;
import cn.fengtang.chatglm.session.OpenAiSession;
import lombok.extern.slf4j.Slf4j;
import okhttp3.sse.EventSource;
import okhttp3.sse.EventSourceListener;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.test.context.junit4.SpringRunner;

import javax.annotation.Resource;
import java.util.Collections;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CountDownLatch;

@Slf4j
@RunWith(SpringRunner.class)
@SpringBootTest
public class ApiTest {
    @Resource
    private OpenAiSession openAiSession;

    @Test
    public void test_chat_completions_stream() throws Exception {
        //1.创建参数
        ChatCompletionRequest chatCompletion = ChatCompletionRequest
                .builder()
                .messages(Collections.singletonList(ChatCompletionRequest.Prompt.builder().role(Role.user.getCode()).content("小傅哥是谁").build()))
                .model(Model.GLM_4V)
                .stream(true)
                .isCompatible(true)
                .build();
        // 2. 发起请求
        EventSource eventSource = openAiSession.completions(chatCompletion, new EventSourceListener() {
            @Override
            public void onEvent(EventSource eventSource, String id, String type, String data) {
                log.info("测试结果：{}", data);
            }
        });
        // 等待
        new CountDownLatch(1).await();
    }
}
